import torch
from netualbuild import NeuralNetwork
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda
import torch.nn as nn

device=(
  "cuda:0"
  if torch.cuda.is_available()
  else "mps"
  if torch.backends.mps.is_available()
  else "cpu"
)

print(f"device: {device}")

model = NeuralNetwork().to(device)
print(model)
torch.save(model,'NeuralNetwork.pth')
model = torch.load('NeuralNetwork.pth',weights_only=False)

labels_map = {
    0: "T-Shirt",
    1: "Trouser",    
    2: "Pullover",    
    3: "Dress",    
    4: "Coat",    
    5: "Sandal",    
    6: "Shirt",    
    7: "Sneaker",    
    8: "Bag",    
    9: "Ankle Boot",
}

training_data = datasets.FashionMNIST(root="data",train=True,download=True,transform=ToTensor())
test_data = datasets.FashionMNIST(root="data", train=False, download=True,transform=ToTensor())

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

learning_rate = 1e-3
batch_size = 64
loss_fn = nn.CrossEntropyLoss()
optmizer = torch.optim.Adam(model.parameters(),lr=learning_rate)
# optmizer = torch.optim.SGD(model.parameters(),lr=learning_rate)

for epochindex in range(0,2):
  model.train()
  for batchindex,(image,label) in enumerate(train_dataloader):
    print(f"image shape:{image.shape}")
    image = image.to(device)
    label = label.to(device)
    pred = model(image)
    print(f"pred shape {pred.shape}")
    print(f"label shape {label.shape}")
    loss = loss_fn(pred,label)
    # if batchindex % 100 == 0:
      # print(f"batchindex:{batchindex} loss:{loss.item()}\n")
    loss.backward()
    optmizer.step()
    optmizer.zero_grad()
  model.eval()
  batchnum = len(test_dataloader)
  datasize = len(test_data)
  test_loss = 0
  test_correct = 0
  with torch.no_grad():
    for batchindex,(image,label) in enumerate(test_dataloader):
      image = image.to(device)
      label = label.to(device)
      pred = model(image)
      test_loss += loss_fn(pred,label)
      # print(f"pred.shape:{pred.shape}  label.shape:{label.shape}\n")
      test_correct += (pred.argmax(1) == label).type(torch.float).sum().item()
      # print(f"pred:{pred.shape} label:{label.shape}, pred[0].size():{pred.size(0)}")
      if torch.equal(pred.argmax(1), label):
        for index in range(0,pred.size(0)):
          predvalue = pred.argmax(1)
          # print(predvalue,label)
          print(f"pred:{labels_map[predvalue[index].item()]} label:{labels_map[label[index].item()]}")
    test_loss /= batchnum
    test_correct /= datasize
    print(f"accuacy:{(100*test_correct):>0.1f}% avgloss:{test_loss:>8f}\n")